# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: donggu

from experimental.prophet.contrib.variable import last_valid
from experimental.prophet.ops.elemwise_math import log
from experimental.prophet.ops import timeseries
from experimental.prophet.contrib.functional import vectorize


@vectorize
def diff(x, index_window=None, time_window=None):
  assert (index_window is None) or (time_window is None)

  x = last_valid(x)

  if time_window:
    past = timeseries.time_shift(x, window_size=time_window)
  else:
    past = timeseries.shift(x, window_size=index_window or 1)

  return x - past


@vectorize
def diff_ratio(x, index_window=None, time_window=None):
  assert (index_window is None) or (time_window is None)

  x = last_valid(x)

  if time_window:
    past = timeseries.time_shift(x, window_size=time_window)
  else:
    past = timeseries.shift(x, window_size=index_window or 1)

  return x / past


@vectorize
def log_return(x, index_window=None, time_window=None):
  assert (index_window is None) or (time_window is None)

  return log(diff_ratio(x, index_window, time_window))


__all__ = ['diff', 'diff_ratio', 'log_return']
